# coding = utf-8

'''
查找qiye数据集中适合做可视化的数据
'''

import os
import SimpleITK as sitk
from PIL import Image
import numpy as np

def find_data_and_compute_dice(case_id):
    raw_root = "E:\Dataset\Qiye\DongBeiDaXue2\image_venous"
    liver_root = "E:\Dataset\Qiye\DongBeiDaXue2\liver"
    tumor_root = "F:\Dataset\Liver\qiye\DongBeiDaXue2\lesion"

    raw_root = os.path.join(raw_root, sorted(os.listdir(raw_root))[case_id-50])
    liver_root = os.path.join(liver_root, sorted(os.listdir(liver_root))[case_id-50])
    tumor_root = os.path.join(tumor_root, sorted(os.listdir(tumor_root))[case_id-50])

    raw_data = sitk.GetArrayFromImage(sitk.ReadImage(raw_root))
    liver_data = sitk.GetArrayFromImage(sitk.ReadImage(liver_root))
    tumor_data = sitk.GetArrayFromImage(sitk.ReadImage(tumor_root))


    raw_data[raw_data <= -250] =250
    raw_data[raw_data >= 200] = 200

    ours_root = "F:\predict\qiye\ours"
    hdenseunet_root = "F:\predict\qiye\hdenseunet"
    munet_root = "F:\predict\qiye\munet"
    unet_root = "F:\predict\qiye\\unet"

    ours_root = os.path.join(ours_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    hdenseunet_root = os.path.join(hdenseunet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    munet_root = os.path.join(munet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    unet_root = os.path.join(unet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))


    index = 0
    for i in range(liver_data.shape[0]):
        if tumor_data[i].sum() > 0:
            ours_image = os.path.join(ours_root, "{}.png".format(str(index).zfill(3)))
            hdenseunet_image = os.path.join(hdenseunet_root, "{}.png".format(str(index).zfill(3)))
            munet_image = os.path.join(munet_root, "{}.png".format(str(index).zfill(3)))
            unet_image = os.path.join(unet_root, "{}.png".format(str(index).zfill(3)))

            if os.path.exists(ours_image):
                ours = Image.open(ours_image).convert("L")
                ours = np.array(ours)
                ours[ours > 0] = 1
            else:
                ours = np.zeros(tumor_data[i].shape)
            dice_ours = float(2*(ours*tumor_data[i]).sum()) / float(ours.sum() + tumor_data[i].sum())

            if os.path.exists(hdenseunet_image):
                hdenseunet = Image.open(hdenseunet_image).convert("L")
                hdenseunet = np.array(hdenseunet)
                hdenseunet[hdenseunet > 0] = 1
            else:
                hdenseunet = np.zeros(tumor_data[i].shape)
            dice_hdenseunet = float(2 * (hdenseunet * tumor_data[i]).sum()) / float(hdenseunet.sum() + tumor_data[i].sum())

            if os.path.exists(munet_image):
                munet = Image.open(munet_image).convert("L")
                munet = np.array(munet)
                munet[munet > 0] = 1
            else:
                munet = np.zeros(tumor_data[i].shape)
            dice_munet = float(2 * (munet * tumor_data[i]).sum()) / float(munet.sum() + tumor_data[i].sum())

            if os.path.exists(unet_image):
                unet = Image.open(unet_image).convert("L")
                unet = np.array(unet)
                unet[unet > 0] = 1
            else:
                unet = np.zeros(tumor_data[i].shape)
            dice_unet = float(2 * (unet * tumor_data[i]).sum()) / float(unet.sum() + tumor_data[i].sum())

            print(index, round(dice_ours, 2), round(dice_hdenseunet, 2), round(dice_munet, 2), round(dice_unet, 2))


        if liver_data[i].sum() > 0:
            index += 1






if __name__ == '__main__':
    find_data_and_compute_dice(case_id=62)